package com.catmiao.spark.sql.test

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
 * @title: Spark01_sql_test01
 * @projectName spark_study
 * @description: TODO
 * @author ChengMiao
 * @date 2024/3/20 10:44
 */
object Spark03_sql_UDAF {

  def main(args: Array[String]): Unit = {
    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("spark_sql_01_test")

    // 创建 Spark Session对象
    val sparkSession: SparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
    // 涉及转换操作，需要引入转换规则

    val df: DataFrame = sparkSession.read.json("datas/sql/user.json")


    df.createOrReplaceTempView("user")
    // 内置函数 拼接常量前缀
    //    sparkSession.sql("select CONCAT('Name:',username),age from user").show
    // 自定义函数 添加前缀 注册自定义函数
    sparkSession.udf.register("ageAvg", new MyAvgUDAF)
    sparkSession.sql("select ageAvg(age) from user").show




    // 关闭环境
    sparkSession.close()
  }


  /**
   * 自定义聚合函数类 计算年龄的平均值
   * 1. 继承 UserDefinedAggregateFunction
   */
  class MyAvgUDAF extends UserDefinedAggregateFunction {

    /**
     * 输入数据的结构
     *
     * @return
     */
    override def inputSchema: StructType = {
      StructType(
        Array(
          StructField("age", LongType)
        )
      )
    }

    /**
     * 缓冲区数据的结构
     *
     * @return
     */
    override def bufferSchema: StructType = {

      StructType(
        Array(
          StructField("total", LongType), // 总和
          StructField("count", LongType) // 计数

        )
      )
    }


    /**
     * 函数计算结果的数据类型 Out
     */
    override def dataType: DataType = LongType

    /**
     * 函数稳定性
     */
    override def deterministic: Boolean = true

    /**
     * 缓冲区初始化
     */
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
//      buffer(0) = 0L
//      buffer(1) = 0L

      buffer.update(0,0L)
      buffer.update(1,0L)

    }

    /**
     * 根据输入值更新缓冲区数据的方法
     */
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      buffer.update(0,buffer.getLong(0) + input.getLong(0) )

      buffer.update(1,buffer.getLong(1)+1)
    }

    /**
     * 缓冲区数据合并
     */
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {

      buffer1.update(0,buffer1.getLong(0) + buffer2.getLong(0))
      buffer1.update(1,buffer1.getLong(1) + buffer2.getLong(1))
    }


    /**
     * 计算，聚合操作 返回结果
     */
    override def evaluate(buffer: Row): Any = {

      buffer.getLong(0) / buffer.getLong(1)
    }
  }

}

